import os
import re
import jieba

class FileLoader(object):

    def __init__(self,
                 root_path:str,
                 stopwords_path:str,
                 years:tuple=None,
                 max_no:int=None):
        self.root_path = root_path
        self.stopwords_path = stopwords_path
        self.years = years
        self.max_no = max_no
        self.stopwords = set()  # 停用词

    def _load_stopwords(self):
        """
        加载中文停用词表
        """
        print("加载停用词表...")
        with open(self.stopwords_path, 'r', encoding='utf-8') as f:
            for line in f:
                self.stopwords.add(line.strip())
        print(f"停用词表加载完成，共 {len(self.stopwords)} 个停用词。")

    def _preprocess_text(self, text:str):
        """
        预处理文本：去除标点符号、数字、停用词等，并进行分词。
        """
        text = re.sub(r'[^\w\s]', '', text)  # 去除标点符号
        text = re.sub(r'\d+', '', text)  # 去除数字
        text = re.sub(r'\s+', ' ', text)  # 去除换行符 \n 和其他空白字符（如 \r、\t 等）
        words = jieba.lcut(text)  # 分词
        words = [word for word in words if word not in self.stopwords and len(word) > 1 and word.strip() != '']  # 去除停用词
        return words

    def _process_file(self, file_path:str):
        """
        加载并预处理单个文件。
        """
        words = []
        with open(file_path, 'r', encoding='utf-8') as file:
            text = file.read()
            words = self._preprocess_text(text)
        return words

    def load_annual_reports(self):
        """
        加载公司年报文本数据
        公司：年：词数据
        """
        self._load_stopwords()
        documents = {}
        total_year_reports_count = 0
        for year_dir_name in os.listdir(self.root_path):
            year = year_dir_name.split("_")[0]
            if self.years is not None and year not in self.years:
                continue
            year_dir_path = os.path.join(self.root_path, year_dir_name)
            cur_year_reports_count = 0
            for company_file_name in os.listdir(year_dir_path):
                split = company_file_name.split("_")
                no = int(split[0])
                if self.max_no is not None and no > self.max_no:
                    break
                company_name = split[2]
                company_file_path = os.path.join(year_dir_path, company_file_name)
                company_annual_report_content = self._process_file(company_file_path)
                if not documents.__contains__(company_name):
                    documents[company_name] = {}
                documents[company_name][year] = company_annual_report_content
                cur_year_reports_count += 1
                total_year_reports_count += 1
                print(f"【{year}】年公司【{company_name}】年报词数据加载完成，"
                    f"共加载 {len(company_annual_report_content)} 个词，"
                    f"本年共加载年报【{cur_year_reports_count}】份，"
                    f"所有年份共加载年报【{total_year_reports_count}】份")
        return documents